#include "obstacle_detector/obstacle_detector.h"
#include "cuda_runtime.h"

namespace obstacle_detector
{
__constant__ float transform_gpu_[12*2];

__global__ void transform_points(int16_t *points, int point_num, float *transformed_points,
	int point_num1, int point_num2)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=point_num) return;
	int lidar_id=threadid<point_num1 ? 0 : 1;
	int16_t origin_x_int=points[threadid*4+0];
	int16_t origin_y_int=points[threadid*4+1];
	float origin_x=points[threadid*4+0]*1.0/256;
	float origin_y=points[threadid*4+1]*1.0/256;
	float origin_z=points[threadid*4+2]*1.0/256;
	float *transform=transform_gpu_+12*lidar_id;
	float x=origin_x*transform[0]+origin_y*transform[1]+origin_z*transform[2]+transform[3];
	float y=origin_x*transform[4]+origin_y*transform[5]+origin_z*transform[6]+transform[7];
	float z=origin_x*transform[8]+origin_y*transform[9]+origin_z*transform[10]+transform[11];
	// transformed_points[threadid*5+0]=static_cast<int16_t>(x*256);
	// transformed_points[threadid*5+1]=static_cast<int16_t>(y*256);
	// transformed_points[threadid*5+2]=static_cast<int16_t>(z*256);
	transformed_points[threadid*5+0]=x;
	transformed_points[threadid*5+1]=y;
	transformed_points[threadid*5+2]=z;
	transformed_points[threadid*5+3]=points[threadid*4+3];
	transformed_points[threadid*5+4]=0;
}

ObstacleDetector::ObstacleDetector(std::string calib_folder,std::string weight_folder,
	bool debug):
debug_(debug)
{
	map_init_=false;
	vehicleparam_init_=false;
	filtered_points_num_=0;

	float all_transform[12*2];
	read_calib(calib_folder,all_transform);
	cudaMemcpyToSymbol(transform_gpu_, all_transform, sizeof(float) * 24 );


	std::string precision = "fp16";
	std::string lidar_onnx =weight_folder+"lidar.backbone.xyz.onnx";
	std::string head_engien = weight_folder+"/head.bbox.plan";
	std::string fuser_engine = weight_folder+"/fuser.plan";
	float threshold = 0.3; // threshold
	detector_ = bevfusion::create_core_lidar_only(precision,lidar_onnx,head_engien,fuser_engine,threshold);

	pc_conversion_=std::make_unique<pc_conversion::PCConversion>(200000, true);
	pointcloud_filter_=std::make_unique<PointCloudFilter>();
	map_processor_=std::make_unique<MapProcessor>(debug_);
	bbox_processor_=std::make_unique<BBoxPorcessor>(false);
	trajectory_predictor_=std::make_unique<TrajectoryPredictor>(debug_);
	collision_detector_=std::make_unique<CollisionDetector>(debug_);

	cudaMalloc((void **)&origin_points_gpu_,sizeof(int16_t)*200000*4);
	cudaMalloc((void **)&filtered_points_gpu_,sizeof(int16_t)*200000*4);
	cudaMalloc((void **)&points_order_gpu_,sizeof(uint16_t)*200000*2);
	cudaMalloc((void **)&points_gpu_,sizeof(float)*200000*5);
	cudaMalloc((void **)&timestamp_gpu_,sizeof(float)*200000);
	cudaMalloc((void **)&point_road_index_gpu_,sizeof(uint8_t)*200000);
	if (debug_)
	{
		debug_display_=std::make_unique<DebugDisplay>("obstacle_detector");
	}
}

void ObstacleDetector::set_filter_pcdata(const uint8_t *data,int point_num)
{
	pc_conversion_->convert2array_gpu(data,filtered_points_gpu_,points_order_gpu_,point_num);
	filtered_points_num_=point_num;
	// if (debug_)
	// {
	// 	// printf ("filtered_points_num: %d\n",point_num);
	// 	int16_t *points_debug=new int16_t [point_num*4];
	// 	cudaMemcpy(points_debug,filtered_points_gpu_,sizeof(int16_t)*point_num*4,
	// 		cudaMemcpyDeviceToHost);
	// 	debug_display_->set_showinfo(points_debug,point_num);
	// 	delete [] points_debug;
	// }

}

const std::vector<Object>& ObstacleDetector::get_obstacles()
{
	return lst_objects_;
}

void ObstacleDetector::set_pcdata(const uint8_t *data1, const uint8_t *data2, int point_num1,
	int point_num2,double stamp)
{
	pc_conversion_->convert2array_gpu(data1,origin_points_gpu_,point_num1, timestamp_gpu_, stamp);
	pc_conversion_->convert2array_gpu(data2,origin_points_gpu_+point_num1*4,point_num2,
		timestamp_gpu_+point_num1, stamp);
	points_num_=point_num1+point_num2;

	int grid_num=(points_num_+1024-1)/1024;
	transform_points<<<grid_num,1024>>>(origin_points_gpu_,points_num_,points_gpu_,
		point_num1,point_num2);
	// pointcloud_filter_->remove_ground(points_gpu_,points_num_);

	// if (debug_)
	// {
	// 	float *points_debug=new float [points_num_*5];
	// 	cudaMemcpy(points_debug,points_gpu_,sizeof(float)*points_num_*5,
	// 		cudaMemcpyDeviceToHost);
	// 	debug_display_->set_showinfo(points_debug,points_num_);
	// 	delete [] points_debug;
	// }

}

float ObstacleDetector::check_collision(module_base::Location &loc)
{
	loc_=loc;
	cuda_base::CUDA_CHECK_RETURN(cudaGetLastError());
	printf ("points_num_ num: %d\n",points_num_);
	auto bboxes = detector_->forward(points_gpu_, points_num_);
	cuda_base::CUDA_CHECK_RETURN(cudaGetLastError());
	printf ("bboxes num: %d\n",bboxes.size());
	std::vector<Object> all_objects;
	for (int i=0;i<bboxes.size();i++)
	{
		auto bb=bboxes[i];
		if (bb.score<0.3) continue;

		// 0- car
		// 1- truck_head
		// 2- RTG
		// 3- tray_w_container
		// 4- street_lamp
		// 5- QC
		// 6- lock_station
		// 7- other_vehicle
		// 8- pedestrian
		// 9- tray_wo_container
		// printf ("label: %d\n",bb.id);
		Eigen::Vector3f center(bb.position.x,bb.position.y,bb.position.z+0.5*bb.size.h);
		Eigen::Vector3f object_dim(bb.size.l,bb.size.w,bb.size.h);
		float yaw=-3.1415926/2-bb.z_rotation;
		Eigen::MatrixXf bbox = generate_bbox(center,object_dim,-1*yaw);
		std::string id=std::to_string(int64_t(loc.stamp*1000))+std::to_string(i)+std::to_string(bb.id);
		all_objects.push_back(Object(id,bb.id,bbox,loc));
	}

	// if (debug_)
	// {
	// 	debug_display_->show3Dbbox(all_objects);
	// }

	if (!map_init_ || filtered_points_num_==0 || !vehicleparam_init_)
	{
		printf ("init fail:  %d,%d,%d\n",map_init_,filtered_points_num_,
			vehicleparam_init_);
		return -1;
	}

	// printf ("filtered_points_num: %d\n",filtered_points_num_);
	// int16_t *points_save=new int16_t [filtered_points_num_*4];
	// cudaMemcpy(points_save,filtered_points_gpu_,sizeof(int16_t)*filtered_points_num_*4,
	// 	cudaMemcpyDeviceToHost);
	// uint16_t *points_order_save=new uint16_t [filtered_points_num_*2];
	// cudaMemcpy(points_order_save,points_order_gpu_,sizeof(uint16_t)*filtered_points_num_*2,
	// 	cudaMemcpyDeviceToHost);
	// save_points(points_save,points_order_save,filtered_points_num_,loc);
	// delete [] points_save;
	// delete [] points_order_save;
	// save_bbox(all_objects,"bbox.txt");
	// save_bbox(lst_objects_,"lst_bbox.txt");

	// filtered_points_num_=17918;
	// int16_t *points_save=new int16_t [filtered_points_num_*4];
	// uint16_t *points_order_save=new uint16_t [filtered_points_num_*2];
	// read_points(points_save,points_order_save,filtered_points_num_,loc);
	// cudaMemcpy(filtered_points_gpu_,points_save,sizeof(int16_t)*filtered_points_num_*4,
	// 	cudaMemcpyHostToDevice);
	// cudaMemcpy(points_order_gpu_,points_order_save,sizeof(uint16_t)*filtered_points_num_*2,
	// 	cudaMemcpyHostToDevice);
	// delete [] points_save;
	// delete [] points_order_save;
	// read_bbox(all_objects,"bbox.txt");
	// read_bbox(lst_objects_,"lst_bbox.txt");

	// if (debug_)
	// {
	// 	debug_display_->show3Dbbox_compare(all_objects,lst_objects_);
	// }
	// printf ("filtered_points_num: %d\n",filtered_points_num_,all_objects.size());
	int roads_num=0;
	std::vector<float> roads_yaw; //每条路个航向角
	map_processor_->map_process(loc,filtered_points_gpu_,
		points_order_gpu_,filtered_points_num_,
		point_road_index_gpu_,roads_num);
	// printf ("roads_num: %d\n",roads_num);
	if (filtered_points_num_==0)
	{
		lst_objects_.clear();
		return -1;
	}
	bbox_processor_->bbox_process(loc,filtered_points_gpu_,
		filtered_points_num_,all_objects,lst_objects_,
		point_road_index_gpu_,roads_num);
	map_processor_->get_roads_yaw(roads_yaw);
	trajectory_predictor_->get_predicted_trajectory(all_objects,
		roads_yaw);
	float collision_index=collision_detector_->check_collision(loc,all_objects);

	// if (debug_)
	// {
	// 	// printf ("filtered_points_num: %d\n",point_num);
	// 	int16_t *points_debug=new int16_t [filtered_points_num_*4];
	// 	cudaMemcpy(points_debug,filtered_points_gpu_,sizeof(int16_t)*filtered_points_num_*4,
	// 		cudaMemcpyDeviceToHost);
	// 	debug_display_->set_showinfo(points_debug,filtered_points_num_);
	// 	delete [] points_debug;
	// }

	if (debug_)
	{
		debug_display_->show3Dbbox_compare(all_objects,lst_objects_);
		int16_t *points_debug=new int16_t [filtered_points_num_*4];
		cudaMemcpy(points_debug,filtered_points_gpu_,sizeof(int16_t)*filtered_points_num_*4,
			cudaMemcpyDeviceToHost);
		float *points_float=new float [filtered_points_num_*4];
		for (int i=0;i<filtered_points_num_;i++)
		{
			float x=points_debug[i*4+0]*1.0/256;
			float y=points_debug[i*4+1]*1.0/256;
			float new_x=x*cos(-1*loc.yaw)+y*sin(-1*loc.yaw);
			float new_y=-1*x*sin(-1*loc.yaw)+y*cos(-1*loc.yaw);
			points_float[i*4+0]=new_x+loc.x;
			points_float[i*4+1]=new_y+loc.y;
			points_float[i*4+2]=points_debug[i*4+2]*1.0/256;

		}
		// save_points(points_float,filtered_points_num_,loc.stamp);
		debug_display_->set_showinfo_float(points_float,filtered_points_num_);
		delete [] points_debug;
		delete [] points_float;
	}

	lst_objects_=all_objects;
	return collision_index;

}

void ObstacleDetector::set_control_sequence(const float *control_sequence)
{
	if (!vehicleparam_init_)
	{
		return;
	}

	float state[51*4];
	motion_prediction(control_sequence,state);
	collision_detector_->set_state(state);
}

void ObstacleDetector::motion_prediction(const float* control_sequence,
	float *state)
{
	state[0]=0;
	state[1]=0;
	state[2]=control_sequence[0];
	state[3]=0;
	float dt=0.2;
	float WB=vehicle_param_.wheel_base;

	for (int i=1;i<(DT+1);i++)
	{
		float v=state[(i-1)*4+2];
		float yaw=state[(i-1)*4+3];

		state[i*4+0]=state[(i-1)*4+0]+v*cos(yaw)*dt;
		state[i*4+1]=state[(i-1)*4+1]+v*sin(yaw)*dt;
		state[i*4+2]=state[(i-1)*4+2]+control_sequence[i*3+0]*dt;
		state[i*4+3]=state[(i-1)*4+3] + v/WB*tan(control_sequence[i*3+1])*dt;
		state[i*4+3]=math::NormalizeAngle(state[i*4+3]);
	}
	if (debug_)
	{
		float global_states[51*4];
		for (int i=0;i<51;i++)
		{
			float x=state[i*4+0];
			float y=state[i*4+1];
			float new_x=x*cos(-1*loc_.yaw)+y*sin(-1*loc_.yaw)+loc_.x;
			float new_y=-1*x*sin(-1*loc_.yaw)+y*cos(-1*loc_.yaw)+loc_.y;
			global_states[i*4+0]=new_x;
			global_states[i*4+1]=new_y;
			global_states[i*4+2]=state[i*4+2];
			global_states[i*4+3]=state[i*4+3]+loc_.yaw;
		}
		debug_display_->set_state(global_states);
	}
}


void ObstacleDetector::set_vechile_param(const module_base::VehicleParam &vehicle_param)
{
	vehicle_param_=vehicle_param;
	if (debug_)
	{
		debug_display_->set_vehicle_param(vehicle_param_);
	}
	collision_detector_->set_vehicle_param(vehicle_param_);
	vehicleparam_init_=true;
}

void ObstacleDetector::set_map(std::string map_path)
{
	map_processor_->set_map(map_path);
	map_init_=true;
	if (debug_)
	{
		size_t pos = map_path.find(".");
		std::string new_path = map_path.substr(0, pos)+"_structure.osm";
		printf ("new path: %s\n",new_path.c_str());
		lanelet::projection::UtmProjector projector(lanelet::Origin({49, 8.4}));
		lanelet::LaneletMapPtr map = lanelet::load(new_path, projector);
		std::vector<Eigen::MatrixXf> all_road_areas;
		for (auto it=map->laneletLayer.begin();it!=map->laneletLayer.end();it++)
		{
			lanelet::Optional<bool> is_straight = it->attributes()["is_straight"].asBool();
			if (!is_straight || !(*is_straight))
			{
				continue;
			}
			
			int left_size=it->leftBound().size()-1;
			int right_size=it->rightBound().size()-1;
			Eigen::MatrixXf area_vertexs(2,4);
			area_vertexs(0,0)=it->leftBound()[0].x();
			area_vertexs(1,0)=it->leftBound()[0].y();
			area_vertexs(0,1)=it->leftBound()[left_size].x();
			area_vertexs(1,1)=it->leftBound()[left_size].y();
			area_vertexs(0,2)=it->rightBound()[right_size].x();
			area_vertexs(1,2)=it->rightBound()[right_size].y();
			area_vertexs(0,3)=it->rightBound()[0].x();
			area_vertexs(1,3)=it->rightBound()[0].y();
			std::cout<<area_vertexs<<std::endl;
			printf ("vertexs size: %d,%d\n",area_vertexs.rows(),area_vertexs.cols());
			all_road_areas.push_back(area_vertexs);
		}
		printf("all_road_areas %d \n",all_road_areas.size());

		debug_display_->set_road_area(all_road_areas);

	}
}

}//namespace